[None][feat] MuseSpark GB300: port the decode kernel stack, switches become dispatch levers by Tabrizian · Pull Request #21953 · NVIDIA/TensorRT-LLM

Tabrizian · 2026-11-10T01:50:14Z

MuseSpark Open GB300: the decode kernel stack on ToT, as levers

Ports the MuseSpark GB300 decode kernel stack onto feat/musespark-open-agentx
(aaba1deabcaf815fbe291bd4b349c425f698c097) and replaces the environment switches it carried with
dispatch levers. 50 files, +13,401 / −304, in 8 commits: the port, the TP8 defaults, and a second
round covering the remaining decode patches as six further levers, all default off.

Thirteen levers are on by default at TP8; everything else is knob-only. At TP4 and every other size
the defaults are unchanged — those compositions either do not serve these paths or have not been tuned
for them.

Environment variables. 21 musespark switches became 14 + 5 dispatch levers. grep -rnE 'os\.environ|getenv'
over the sm103 target now returns exactly two lines — levers.py's own snapshot_of(os.environ)
and a cuBLAS workspace check — neither a feature switch. sm103 had no lever_settings.py; this adds
one modeled on the sm107 sibling.

Resolution is proved by A/B rather than asserted. A verifier rebuilds FAST_PATHS, DEPENDS_ON
and both compositions by AST from each tree, resolves with an empty environment, and compares
resolved path sets. On the port commit they are identical to the base at all ten deployments — 44
paths at TEP8, 54/53 at TEP4 — and each lever is then shown to turn on by its own knob. The defaults
commit moves only TP8: every other deployment resolves exactly as before.

A defect this port introduced, and fixes

weights._ring_buddy_row_cap called max_rows_for(compute_capability, tp_size) — the signature the
serving image had — while the ported max_rows_for takes a single ranks. moe-expert-balance/1
wrote that call site; the arity only became reachable once experts-early-start landed
max_rows_for itself. That function is called unconditionally from the replicated_experts install,
so every TEP8 deployment would have raised TypeError at load_weights. The call now passes
ring_size(mapping), which is what install() itself passes.

Nothing caught it earlier because nothing had run on a GPU, and the lever verifier resolves tables
without touching the install path.

Behaviour change worth calling out

All 21 retired names are in levers.RETIRED, so a deployment still exporting one raises and names
its replacement knob rather than running with the switch silently ignored. That is the mechanism's
point, but it means any config exporting them fails to start until updated. Checked: no recipe under
agentx/ or agent-perf/ in the benchmarking repo exports any of them.

What is measured

The branch with every lever off is close to ToT, measured cross-holder. Same driver, same image
lineage, GB300 TEP8, decode gpuForwardTimeMS medians:

batch	gen	ToT `aaba1deabca`	this branch, levers off	delta	n
3	3	11.925 ms	11.666 ms	−2.17 %	54
4	3	11.817 ms	11.773 ms	−0.37 %	7
4	4	12.283 ms	12.258 ms	−0.20 %	47

Caveat, and it is not small: the two arms did not share an allocation or a node set, and they ran
several hours apart. The ~2 % floor quoted elsewhere comes from
within-run repeats (gen = 3 measured twice in one tree), which does not bound node-to-node or
hour-to-hour variation — and this project has on record a same-config point spanning 78-116k
tok/s/GPU by node set and time. So these deltas are consistent with no change, but they do not
establish it to 2 %. A same-holder A/B, both trees in one allocation, would.

Read that as a correctness result, not a performance one. It says the port did not disturb the
default path on hardware, which the static A/B could not show.

Numerical verification

No accuracy result is claimed for these levers. Several of the layers are non-bitwise, and a lever
can reproduce the base's output at one row count and not at another — this stack has already produced
one such case, where a kernel assumed a wide table was armed at a row band that this fork serves from
cuBLAS instead. So a gate has to run at the batch sizes these paths are written for, and that has not
been done. The project's numerical-verification gate must pass before any point using them counts.

Two levers are deliberately not default-on. drafter_fused_prep resolves off by dependency on this
composition. ring_buddy_stats is observability only and is mutually exclusive with
tep8_experts_in_kernel_placement.

TRTLLM_MULTI_INSTANCE_CONTEXT_FMHA is not a lever and is not enabled here. A decode-step speedup
measured for this stack was taken with that switch also on, so the levers' own share of it is not
separated and is not claimed.

No performance claim from the original patch READMEs carries over. Their numbers were measured on
the pre-split tree against a wide_projection_gathers that exists in no public commit. They are not
reproduced here and are not claimed.

Known limits, recorded rather than papered over

tep8_wide_direct_transports is served at TEP4 despite its name. Its call site needs
wide_projection_gathers, which no sm103 TEP8 composition serves, so the env var was already a dead
no-op at TEP8. It is now declared with that dependency, so a knob set at TEP8 raises instead of being
ignored. Making it genuinely reachable is a composition change, not a port.
ring_buddy_stats lost sub-second intervals: Param parses integers and the env accepted 0.5.
tep8_ring_buddy_wide__MAX_ROWS keeps the env's full 1..256 range, so it can still lower the cap
below the built-in 16 despite the name.
The placement/stats exclusion stays the runtime guard _stats is None in _State.current rather
than a Rule.conflicts, which would turn today's "both set, stats win" into a resolve-time error —
a behaviour change beyond a port.
state.wide is new in this base and has not been exercised together with
tep8_experts_in_kernel_placement. The gates look compatible — both need nvfp4, two windows,
balance and 8 ranks — but that is an argument, not a measurement.
early_trigger is not wired into the new wide-assign kernel: early_trigger implies ≤32 rows → ≤256
slots, while the wide path needs >512, so the branches are disjoint today. If BATCH_MAX_ROWS rises
above 64 that argument lapses and the wide launch's PDL=False would silently drop the early
trigger — slower, not incorrect.

Defects this port introduced, found by running it

Static checks passed on all three; only execution surfaced them.

weights._ring_buddy_row_cap called max_rows_for with two arguments against a one-argument
signature. On the unconditional replicated_experts install path, so every TP8 deployment would
have raised TypeError at load_weights.
moe_prologue declared _wide_direct: bool as a class attribute, shadowing
projection_gathers' transport object through the MRO. Every 8-row decode step raised
AttributeError at warmup, in a path served by default at TP8.
pool_into_route called _wide_router_shard with four arguments against five. Latent, behind a
default-off lever.

…become dispatch levers Ports the MuseSpark GB300 decode kernel stack onto feat/musespark-open-agentx and replaces the environment switches it carried with dispatch levers: staged router GEMM, ring-buddy cap and wide balance, router shard, drafter fast path, wide direct transports and pool, persistent experts, experts early start and interleaved fetch, in-kernel placement, finalize entry trigger, the SWA cluster block, the widened TP8 finalize transport, batch split-K, and the multi-instance context FMHA. EVERY LEVER IS DEFAULT OFF. This is the port; it changes no behaviour until a lever is knobbed on. Proved rather than asserted: a verifier rebuilds FAST_PATHS, DEPENDS_ON and both compositions by AST and resolves them with an EMPTY environment -- the resolved path sets are identical to the base at all ten deployments (44 paths at TEP8, 54/53 at TEP4), and each lever is then shown to turn on by its own knob. THREE DEFECTS FIXED, each found by running the code after static checks passed: 1. weights._ring_buddy_row_cap called max_rows_for with two arguments against a one-argument signature, on the unconditional replicated_experts install path. Every TP8 deployment would have raised TypeError at load_weights. 2. moe_prologue declared _wide_direct: bool as a class attribute, shadowing projection_gathers' transport object (default None) through the MRO. Every 8-row decode step raised AttributeError("'bool' object has no attribute 'is_armed'") at warmup, inside attention_output_into_hop, which is served by default at TP8. 3. pool_into_route called _wide_router_shard with four arguments against five. Latent, behind a default-off lever. THIRD-PARTY LICENSING. multi_instance_context_kernel.py derives from the CUTLASS CuTe DSL Blackwell FMHA example (v4.8.0) and declares that derivation with the full BSD-3-Clause notice in-file, matching 51 existing files in this repository and the sm107 sibling added in 736cfdc. CUTLASS is already attributed in ATTRIBUTIONS-Python.md and both per-arch CPP files, so no attribution file changes. MEASURED. With every lever off this branch is indistinguishable from its base on GB300 TEP8 decode (11.67 vs 11.93 ms at bs3, 12.26 vs 12.28 at bs4), which is a correctness result -- the port does not disturb the default path -- and says nothing about what the levers are worth. Non-bitwise levers still need their acceptance, drift and step-time gates before any of them ships on by default. Signed-off-by: Iman Tabrizian <itabrizian@nvidia.com>

Thirteen levers of the sm103 decode stack move from knob-only to on at TP8, via a _tep8_only rule. At TP4 and every other size they stay knob-only: those compositions either do not serve these paths or have not been tuned for them. What carries the default is acceptance, measured on the SPEED-Bench 80-row coding set (chat-templated, use_rejection_sampling, DFlash d7, concurrency 1, GB300 TEP8): with all thirteen on, the run is token-identical to the all-off stack -- AL 4.6303, the same 3623/6986 accepted/drafted, the same output. Every arm of that gate repeats bitwise, so the result is exact rather than inside a band, and token identity makes it independent of the baseline's absolute level. Two levers were deliberately left knob-only: * drafter_fused_prep resolved off by dependency in BOTH measured arms, so the bitwise-clean result does not cover it. * ring_buddy_stats is measured clean but is observability only, and it is mutually exclusive with tep8_experts_in_kernel_placement. TRTLLM_MULTI_INSTANCE_CONTEXT_FMHA is not a lever and is not enabled by this change. On the same gate it costs 3.79 % AL -- all of the loss the full stack showed. The decode step measurement quoted for this stack (12.23 -> 11.04 ms at bs4) was taken with that switch on, so the levers' own share of it is not separated and is not claimed here. Signed-off-by: Iman Tabrizian <itabrizian@nvidia.com>

Tabrizian · 2026-11-10T05:40:40Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-10T05:48:04Z

PR_Github #83672 [ run ] triggered by Bot. Commit: 0db4fd5 Link to invocation

CI runs pre-commit over the whole tree while a local run only sees staged files, so two repo hooks were failing on this branch from the port commit on. * ruff-format reformats six files. Verified format-only: each one's AST is identical before and after, so nothing changed meaning. * The pinned-memory policy disallows raw pinning in runtime code. The ring-buddy placement-stats accumulator built its host buffer with torch.zeros(..., pin_memory=True); it now goes through tensorrt_llm._utils.maybe_pin_memory(), which is what the policy asks for and is idempotent. That buffer belongs to ring_buddy_stats, which is knob-only, so no default path is affected. Signed-off-by: Iman Tabrizian <itabrizian@nvidia.com>

tensorrt-cicd · 2026-11-10T08:41:17Z

PR_Github #83672 [ run ] completed with state SUCCESS. Commit: 0db4fd5
/LLM/main/L0_MergeRequest_PR pipeline #69181 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Tabrizian · 2026-11-10T14:37:04Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-10T14:45:30Z

PR_Github #83811 [ run ] triggered by Bot. Commit: f6a8a2e Link to invocation

…evers Folds two drafter-tail launch pairs of the sm103 TEP8 decode step, each behind a dispatch lever that is default off: * drafter_conv_epilogue (on drafter_block_conv and drafter_fixed_gemms): on blocks of 9..32 rows the coefficient GEMM computes the input-side convolution in its epilogue and prefetches its operands under the mainloop, over a row-permuted copy of the projection weight (32 MiB per convolution). The weight values are unchanged. * drafter_finish_norm (on drafter_block_decode): on blocks the drafter hops do not take, each hop's finish convolution, residual add and RMSNorm are one launch that releases its dependents at entry. Both are lossless-bitwise and knob-only, so an environment that sets nothing runs exactly the paths it ran before. The two environment switches the work carried are retired by name in levers.RETIRED, so a deployment still exporting one is told which knob replaces it instead of being silently ignored. Served at TEP8 only: these are base fast paths, and the sm103 TEP4 composition lists its own set, so the resolver refuses the knob there rather than serving an unmeasured shape. Proved rather than asserted: a verifier rebuilds FAST_PATHS, DEPENDS_ON and both compositions by AST and resolves them with an EMPTY environment. The resolved path sets are identical to the base at all twelve deployments checked, and each new lever then turns on by its own knob. Run on GB300 (2 nodes, TEP8 holder): all 16 touched modules import, including this commit's two new CuTe DSL files; the real lever table resolves to 44 TEP8 paths with an EMPTY environment, matching the AST verifier; conv_coef_gemm compiles and runs with finite outputs at 16 rows. The model itself was not booted. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 34d9a4d0b161a427d0e893b73906fbcb695c550d)

…ttention-block split-K Three more default-off dispatch levers on the sm103 TEP8 decode step: * drafter_multi_request_cells (on drafter_fixed_gemms): the drafter GEMMs of 16, 24 and 32-row blocks run measured TP8 split-K cells. The qkv, gate_up and coefficient cells keep their table row's split_k and are bitwise that row; the o_proj and down_proj cells are single K-slice GEMMs that reproduced the library call's bits on every checked shape, which cannot be proven against a closed library, so the lever is flagged rather than called bitwise. * drafter_silu_epilogue (on drafter_block_decode): gate_up_proj and silu(gate) * up of 9..32-row blocks are one launch over an interleaved copy of the gate_up weights, which costs about 560 MiB per generation rank. * tep8_attn_block_splitk (on dense_gemms): the attention block's QKV and o_proj of 9..32-row decode calls run batch split-K cells, so each stages its weight tile before its dependency wait. o_proj at split_k 1 keeps the library call's summation; QKV at split_k 4 keeps the served cell's. A tactic the drafter table marks OnBatchSplitK is run on this target's own split-K copy, whose tiles and summation match the served kernel at equal split_k. Every lever is knob-only, so an environment that sets nothing runs exactly the paths it ran before; each environment switch this work carried is retired by name in levers.RETIRED. Proved by the AST verifier: rebuilding FAST_PATHS, DEPENDS_ON and both compositions and resolving with an EMPTY environment gives path sets identical to the base at all twelve deployments, and each lever then turns on by its own knob. Run on GB300 (2 nodes, TEP8 holder): all 13 drafter cells validate against the split-K kernel's own bounds; the whole attention-block table passes its checker and is refused at TP4; gate_up_silu_gemm is bitwise the GEMM plus silu_and_mul it replaces at 16, 24 and 32 rows (0 differing elements of 57344, 86016 and 114688), once the reference does its arithmetic in FP32 and rounds once; every attention-block cell compiles and runs, and the o_proj cells at split_k 1 are bitwise cuBLAS at 9, 17 and 25 rows. The model itself was not booted. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 8bb1c39fa67854227cef8030a29737ad7103903d)

… decode rows Raises the ceiling of tep8_wide_balance's MAX_ROWS from 256 to 512 rows. Above 256 rows (32 requests at seven draft tokens) the balanced NVFP4 placement had no band left, so a step past that count fell back where the row-count census, the ring-buddy cap and the placement band all ended at once. One wide ring_buddy_assign program now covers a 4096-slot tile. Its work is per expert -- the census is a histogram of the selected ids, the move decision is one per expert and the gather is one per slot -- so the larger tile costs no [SLOTS, RANKS] tensor. The narrow program keeps its 2048-slot bound, and a call takes the wider one only where wide_serves already admits it: the balanced TP8 placement above the wide program's minimum, with no half window and no static move. The lever's default stays 256, so nothing reaches the new band until the parameter is set, and the lever itself is still default off. The AST verifier resolves the same path sets as the base at all twelve deployments with an EMPTY environment; MAX_ROWS now accepts 288 and 512 and still refuses 128 and 513. Run on GB300 (2 nodes, TEP8 holder): ring_buddy_assign's widened validation accepts 288 and 256 rows on the wide program and still takes 32 on the narrow one, exercised on the unconditional path. The model itself was not booted. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 90a880ae9fe98b908cfd95e8f69cda7ebec4bf52)

…oads The sliding-window block's split and combine kernels already take an EARLY_LOADS build, but nothing asked for it, so the capability was inert. swa_block_single_early_loads (on swa_block, default off) arms it on the one-request block: * the split issues every projection load the program makes after its wait at once -- the query rows, the key block's step K/V rows and, in program (0, h, 0), the step's K/V rows it writes to the cache -- before any prep arithmetic, so the step costs one dependent round trip instead of one per prep. It needs the fused prep, whose operands those loads are, and is passed only when both hold; * the combine loads the partial O with the (max, sum) words, and the gate head before its wait where there is a gate. The same statements on the same values, so the same bits. The blocks-per-program split runs a different kernel that has no such build, so it keeps the served launch. Knob-only and default off; the environment switch this work carried is retired by name in levers.RETIRED. The AST verifier resolves path sets identical to the base at all twelve deployments with an EMPTY environment, and the lever turns on by its own knob at TEP8. Run on GB300 (2 nodes, TEP8 holder): the EARLY_LOADS build this commit arms -- which nothing had compiled before, since the halves were ported without a caller -- compiles, runs, and is bitwise the served build (0 differing elements). The model itself was not booted. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 2506fb61c71ed363e3ca4f14e8b1941c8eba2cbe)

…itwise A batch-4 run on GB300 showed this lever changing the decoded text: against an otherwise identical arm it more than doubled the decode steps for the same token budget (287 -> 641), while the five other levers of this series reproduced the step structure exactly. Greedy decoding, same prompts, one knob apart. The cause is a premise that does not hold at TEP8. The lever replaces the 9..32-row QKV call with a split_k 4 cell, which keeps the served bits only if the served call at those rows is already that cell. It is not here: the path that arms a wide table at those rows, wide_projection_gathers, is TEP4-only on sm103, so the served 9..32-row QKV call is the built-in cuBLAS one, which sums K in a single slice. Four slices instead of one moves the last bits, and the sampled text follows. The o_proj half is unaffected and stays bitwise: its split_k 1 cells reproduce cuBLAS exactly, checked here at 9, 17 and 25 rows. So the lever becomes lossless-numeric, which also means auto needs a gate for it rather than treating it as free. It stays default off, and the module now says where its QKV cells are bitwise and where they are not, instead of claiming every row keeps its bits. Measured alongside: this lever costs 0.265 ms on the bs4 decode step (12.440 -> 12.705 ms, ~270 steps per arm), and is the whole of the 0.195 ms the six levers cost together. The other five sit within 0.021 ms of the baseline. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 28fb09a9b987a12c8cc60f66bf25a58ad6192f88)

Tabrizian · 2026-11-10T14:57:53Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-10T15:05:14Z

PR_Github #83816 [ run ] triggered by Bot. Commit: e333e03 Link to invocation

tensorrt-cicd · 2026-11-10T15:09:38Z

PR_Github #83811 [ run ] completed with state ABORTED. Commit: f6a8a2e

Link to invocation

tensorrt-cicd · 2026-11-10T17:50:53Z

PR_Github #83816 [ run ] completed with state SUCCESS. Commit: e333e03
/LLM/main/L0_MergeRequest_PR pipeline #69308 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…d levers test_a_paged_context_plans_every_fast_path_but_the_window_compaction asserts the exact set of levers that resolve off on a sm103 TP4 paged-context deployment. The decode-stack port adds five more levers, so that set grows from two to seven. All five are correctly off there: four are default off everywhere, and finalize_entry_trigger carries its default at TP8 only. No production code changes; the assertion was enumerating a set the port legitimately extends. Signed-off-by: Iman Tabrizian <itabrizian@nvidia.com>

Tabrizian · 2026-11-10T18:11:52Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-10T18:18:43Z

PR_Github #83861 [ run ] triggered by Bot. Commit: 0b10da0 Link to invocation

github-actions Bot assigned Tabrizian Nov 10, 2026

Tabrizian marked this pull request as ready for review November 10, 2026 05:40

Tabrizian requested review from a team, Wanli-Jiang, WeiHaocheng and litaotju as code owners November 10, 2026 05:40

Tabrizian requested a review from pengbowang-nv November 10, 2026 05:40

Tabrizian added 5 commits November 10, 2026 06:50

Navigation Menu

Repository navigation

[None][feat] MuseSpark GB300: port the decode kernel stack, switches become dispatch levers - #21953

[None][feat] MuseSpark GB300: port the decode kernel stack, switches become dispatch levers#21953
Tabrizian wants to merge 9 commits into
NVIDIA:feat/musespark-open-agentxfrom
Tabrizian:user/itabrizian/musespark-gb300-decode-kernel-stack

Tabrizian commented Nov 10, 2026 •

edited

Loading

Uh oh!

Tabrizian commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

Tabrizian commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

Tabrizian commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

Tabrizian commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Tabrizian commented Nov 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

MuseSpark Open GB300: the decode kernel stack on ToT, as levers

A defect this port introduced, and fixes

Behaviour change worth calling out

What is measured

Numerical verification

Known limits, recorded rather than papered over

Defects this port introduced, found by running it

Uh oh!

Tabrizian commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

Tabrizian commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

Tabrizian commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

Tabrizian commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Tabrizian commented Nov 10, 2026 •

edited

Loading